from py2neo import Node, Relationship, Graph, NodeMatcher, RelationshipMatcher, Subgraph
import global_var

class DataBase:

    def __init__(self):
        # self.graph = Graph("http://localhost:7474", password="123456", name="neo4j")
        self.graph = Graph("http://211.71.15.39:9474", password="python", name="python")
        self.nodes = NodeMatcher(self.graph)
        self.relations = RelationshipMatcher(self.graph)

    # 删除图的所有内容
    def delete_all(self):
        self.graph.delete_all()

    # 查找所有符合给定标签的结点
    def print_nodes(self, tag):
        node_tag = self.nodes.match(tag).all()
        print("%s类的%s个节点如下：" % (tag, len(node_tag)))
        for node in node_tag:
            print(node)

    # 查找所有符合给定标签的关系
    def print_relations(self, tag):
        relation_tag = self.relations.match(tag).all()
        print("%s类的%s个关系如下：" % (tag, len(relation_tag)))
        for node in relation_tag:
            print(node)

    # 通过id查找包
    def match_pkg_by_id(self, id):
        res = self.nodes.match('Package', ID=id)
        return res.first()

    # 通过name查找项目
    def match_proj_by_name(self, name):
        res = self.nodes.match('Project', name=name)
        return res.first()

    # 通过id查找开发者
    def match_contri_by_id(self, id):
        res = self.nodes.match('developer', github_id=id)
        return res.first()

    # 通过url查找仓库
    def match_repository_by_url(self, url):
        res = self.nodes.match('Repository', url=url)
        return res.first()

    # 通过id查找漏洞
    def match_vulnerability_by_id(self, id):
        res = self.nodes.match('Vulnerability', ID=id)
        return res.first()

    def test(self):
        self.graph.create(Node('test', a=1, b=2, c=3))

    # 添加包结点
    def add_pkgs(self, filecount):
        global_var.load_pkg_node(filecount)
        count = 1
        for node in global_var.nodes:
            self.graph.create(Node('Package', ID=node.ID, name=node.name, version=node.version))
            if count % 500 == 0:
                print("already save %s nodes" % count)
            count += 1
        print("file %s save nodes has done!" % filecount)
        global_var.nodes.clear()

    # 添加项目结点
    def add_projs(self):
        global_var.load_proj_node()
        count = 1
        for node in global_var.nodes:
            if self.match_proj_by_name(node.name) is not None:
                continue
            self.graph.create(Node('Project', name=node.name, info=node.info, license=node.license, homepage=node.homepage, repository_url=node.repository_url, stars=node.stars))
            if count % 500 == 0:
                print("already save %s nodes" % count)
            count += 1
        print("save project nodes has done!")
        global_var.nodes.clear()

    # 添加开发者结点
    def add_contributors(self):
        global_var.load_contri_node()
        count = 1
        for node in global_var.nodes:
            if self.match_contri_by_id(node.github_id) is not None:
                continue
            self.graph.create(
                Node('developer', github_id=node.github_id, github_name=node.github_name, name=node.name, homepage=node.homepage,
                     organization=node.organization, email=node.email, address=node.address))
            if count % 500 == 0:
                print("already save %s nodes" % count)
            count += 1
        print("save %s developer nodes has done!" % count)
        global_var.nodes.clear()

    # 添加仓库结点
    def add_repository(self, repo_detail):
        for key in repo_detail:
            info = repo_detail[key]
            if info == {}:  # 没有repo信息的仓库只添加name和url即可
                flag = 0
                for i in range(len(key)-1, 0, -1):
                    if key[i] == '/':
                        flag = i + 1
                        break
                self.graph.create(
                    Node('Repository', name=key[flag:], url=key))
            else:
                self.graph.create(
                    Node('Repository', name=info.get('name'), url=info.get('html_url'), language=info.get('language'),
                         owner=info.get('owner').get('id'), contributor=info.get('contributors_url')))
        print("save %s developer nodes has done!" % len(repo_detail))
        global_var.nodes.clear()

    def save_relation_list(self, relation_list):
        subgraph = Subgraph(relationships=relation_list)
        transaction = self.graph.begin()
        self.graph.create(subgraph)
        transaction.commit()

    def add_commits(self, commits):
        relation_list = []
        for url in commits:
            node1 = self.match_repository_by_url(url)
            if node1 is None:
                continue
            commit_list = commits[url]
            if commit_list == {}:
                continue
            for commit in commit_list:
                try:
                    node2 = Node('commit', ID=commit.get('node_id'), developer_ID=commit.get('committer').get('id'),
                                 commit_message=commit.get('commit').get('message'), committer=commit.get('committer').get('login'))
                    self.graph.create(node2)
                    relation_list.append(Relationship(node2, "commit_update_repo", node1))
                except:
                    pass
            self.save_relation_list(relation_list)
            print("finish %s relations" % len(relation_list))
            relation_list.clear()

    def add_issues(self, issues):
        relation_list = []
        for url in issues:
            node1 = self.match_repository_by_url(url)
            if node1 is None:
                continue
            issue_list = issues[url]
            if issue_list == {}:
                continue
            for issue in issue_list:
                try:
                    node2 = Node('issue', ID=issue.get('id'), developer_ID=issue.get('user').get('id'), API_ID=issue.get('number'),
                                 title=issue.get('title'), body=issue.get('body'), create_time=issue.get('created_at'), update_time=issue.get('updated_at'),
                                 finish_time=issue.get('closed_at'), state=(int)(issue.get('state')=='open'), comments=issue.get('comments_url'))
                    self.graph.create(node2)
                    relation_list.append(Relationship(node1, "issue_in_repo", node2))
                except:
                    pass
            self.save_relation_list(relation_list)
            print("finish %s relations" % len(relation_list))
            relation_list.clear()

    # 添加项目-包关系
    def add_pkg_proj_relation(self, start, end):
        relation_list = []
        print("add pkg-proj relations from file %s to file %s to db" % (start, end - 1))
        for file_count in range(start, end):
            pkg_id = global_var.load_pkg_id(file_count)
            count = 0
            for id in pkg_id:
                pkg = id.split('@')[0]
                node1 = self.match_proj_by_name(pkg)
                node2 = self.match_pkg_by_id(id)
                if node1 is not None and node2 is not None:
                    relation_list.append(Relationship(node1, "include_pkg", node2))
                    relation_list.append(Relationship(node2, "part_of_project", node1))
                    if len(relation_list) >= 1000:  #每1000个关系存一下，用子图的方式效率更高
                        self.save_relation_list(relation_list)
                        relation_list.clear()
                        count += 1000
                        print("finish %s relations" % count)
                else:
                    print("fail to create %s-%s realtions" % (pkg, id))

            if len(relation_list) != 0:
                self.save_relation_list(relation_list)
                count += len(relation_list)
                print("finish %s relations" % count)
                relation_list.clear()

            print("file %s has done!" % file_count)
        print("save proj-pkg relations has done!")

    # 添加项目-开发者关系
    def add_proj_contri_relation(self, start, end):
        relation_list = []
        for file_count in range(start, end):
            contributor_dict = global_var.load_contri_dict(file_count)
            count = 0
            for key in contributor_dict:
                contri_list = contributor_dict[key]
                if isinstance(contri_list, int):
                    continue
                node1 = self.match_proj_by_name(key)
                if node1 is None:
                    continue
                for contri in contri_list:
                    node2 = self.match_contri_by_id(contri['github_id'])
                    if node2 is None:
                        node2 = Node('developer', github_id=contri.get('github_id'), github_name=contri.get('login'), name=contri.get('name'),
                                 homepage=contri.get('blog'), organization=contri.get('company'), email=contri.get('email'), address=contri.get('location'))
                        self.graph.create(node2)
                    relation_list.append(Relationship(node1, "include_developer", node2))
                    relation_list.append(Relationship(node2, "develop", node1))
                    if len(relation_list) >= 1000:  # 每1000个关系存一下，用子图的方式效率更高
                        self.save_relation_list(relation_list)
                        relation_list.clear()
                        count += 1000
                        print("finish %s:%s relations" % (file_count, count))

            if len(relation_list) != 0:
                self.save_relation_list(relation_list)
                count += len(relation_list)
                print("finish %s:%s relations" % (file_count, count))
                relation_list.clear()

            print("file %s has done!" % file_count)

        print("save proj-contri relations has done!")

    # 添加包-包关系
    def add_pkg_pkg_relation(self, start, end):
        relation_list = []
        print("add pkg-pkg relations from file %s to file %s to db" % (start, end-1))
        for file_count in range(start, end):
            pkg_dep_pkg = global_var.load_pkg_dep_pkg(file_count)
            count = 0
            for dep in pkg_dep_pkg:
                node1 = self.match_pkg_by_id(dep.get('pkg_a'))
                node2 = self.match_pkg_by_id(dep.get('pkg_b'))
                if node1 is not None and node2 is not None:
                    relation_list.append(Relationship(node1, "depends_on_pkg", node2))
                    if len(relation_list) >= 1000:  #每1000个关系存一下，用子图的方式效率更高
                        self.save_relation_list(relation_list)
                        relation_list.clear()
                        count += 1000
                        print("finish %s:%s relations" % (file_count, count))
                else:
                    print("fail to create %s-%s realtions" % (dep.get('pkg_a'), dep.get('pkg_b')))

            if len(relation_list) != 0:
                self.save_relation_list(relation_list)
                count += len(relation_list)
                print("finish %s:%s relations" % (file_count, count))
                relation_list.clear()

            print("file %s has done!" % file_count)

        print("save pkg-pkg relations has done!")

    # 添加漏洞信息和漏洞与包结点的关系
    def add_vul_pkg_relation(self):
        relation_list = []
        for file_count in range(1, 58):
            vul_dict = global_var.load_vul_dict(file_count)
            count = 0
            for key in vul_dict:
                node1 = self.match_pkg_by_id(key)
                if node1 is None:
                    continue
                vuls = vul_dict[key]
                if len(vuls) == 0:
                    continue
                for vul in vuls:
                    node2 = self.match_vulnerability_by_id(vul.get('id'))
                    if node2 is None:
                        if len(vul.get('aliases')) > 0:
                            node2 = Node('Vulnerability', ID=vul.get('id'), CVE_ID=vul.get('aliases')[0], info=vul.get('details'), url=vul.get('link'))
                        else:
                            node2 = Node('Vulnerability', ID=vul.get('id'), CVE_ID="", info=vul.get('details'), url=vul.get('link'))
                        self.graph.create(node2)
                    relation_list.append(Relationship(node2, "affect_pkg", node1))

                if len(relation_list) >= 1000:
                    self.save_relation_list(relation_list)
                    count += len(relation_list)
                    print("finish %s:%s relations" % (file_count, count))
                    relation_list.clear()

            if len(relation_list) != 0:
                self.save_relation_list(relation_list)
                count += len(relation_list)
                print("finish %s:%s relations" % (file_count, count))
                relation_list.clear()


            print("file %s has done!" % file_count)

        print("save vul-pkg relations has done!")